import { NodeTypes } from '@xyflow/react';
import React from 'react';
import CustomNodeComponent from './CustomNode';
import InputNode from './InputNode';
import ProcessingNode from './ProcessingNode';
import FeatureSelectionNode from './FeatureSelectionNode';
import MachineLearningNode from './MachineLearningNode';
import VisualizationNode from './VisualizationNode';
import ModelEvaluationNode from './ModelEvaluationNode';
import { NODE_TYPES } from './NodeTypes';

// 定义节点类型映射
export const nodeTypeMapping = {
  custom: (props: any) => {
    const { data } = props;
    if (!data) return React.createElement('div', null, '无效节点');

    // 根据节点类型返回对应组件
    if (['csv_reader', 'excel_reader', 'db_reader'].includes(data.type)) {
      return React.createElement(InputNode, props);
    } else if (['data_cleaning', 'data_standardization', 'outlier_detection', 'one_hot_encoding'].includes(data.type)) {
      return React.createElement(ProcessingNode, props);
    } else if (['feature_selection_filter'].includes(data.type)) {
      return React.createElement(FeatureSelectionNode, props);
    } else if (['logistic_regression', 'random_forest', 'svm'].includes(data.type)) {
      return React.createElement(MachineLearningNode, props);
    } else if (['structured_view', 'bar_chart', 'line_chart', 'scatter_plot', 'pie_chart', 'heatmap', 'model_evaluation'].includes(data.type)) {
      if (data.type === 'model_evaluation') {
        return React.createElement(ModelEvaluationNode, props);
      }
      return React.createElement(VisualizationNode, props);
    }

    return React.createElement('div', null, `未知节点类型: ${data.type}`);
  }
};

// 通过节点类型获取合适的节点组件
export const getNodeComponent = (nodeType: string) => {
  if (nodeType === 'custom' && nodeTypeMapping.custom) {
    return nodeTypeMapping.custom;
  }
  return CustomNodeComponent; // 默认回退到CustomNodeComponent
};

export { default as CustomNodeComponent } from './CustomNode';
export { default as InputNode } from './InputNode';
export { default as ProcessingNode } from './ProcessingNode';
export { default as MachineLearningNode } from './MachineLearningNode';
export { default as VisualizationNode } from './VisualizationNode';
export { default as ModelEvaluationNode } from './ModelEvaluationNode';
export * from './NodeTypes'; 